from .capture_base import CaptureBase
import win32gui, win32ui, win32con
import numpy as np


class Win32Capture(CaptureBase):
    def __init__(self, hwnd=None, width=320, height=320):
        """
        :param hwnd: 目标窗口句柄；None 表示当前前台窗口
        :param width: 截图区域宽度
        :param height: 截图区域高度
        """
        self.hwnd = hwnd or win32gui.GetForegroundWindow()
        self.width = width
        self.height = height
        self.region = None

    def init_in_thread(self):
        """初始化截图区域"""
        self._update_region()

    def _update_region(self):
        """计算窗口中心区域"""
        left, top, right, bottom = win32gui.GetWindowRect(self.hwnd)
        win_w = max(1, right - left)
        win_h = max(1, bottom - top)

        cx = left + win_w // 2
        cy = top + win_h // 2

        self.region = (
            cx - self.width // 2,
            cy - self.height // 2,
            self.width,
            self.height,
        )

    def update_size(self, width, height):
        """热更新截图区域大小"""
        self.width = width
        self.height = height
        self._update_region()
        print(f"🔄 Win32Capture 截图大小已更新: {self.width}x{self.height}")

    def grab_frame(self):
        """抓取一帧 (H, W, 3)，BGR"""
        if self.region is None:
            self._update_region()

        left, top, w, h = self.region

        hwnd_dc = win32gui.GetWindowDC(self.hwnd)
        src_dc = win32ui.CreateDCFromHandle(hwnd_dc)
        mem_dc = src_dc.CreateCompatibleDC()

        bmp = win32ui.CreateBitmap()
        bmp.CreateCompatibleBitmap(src_dc, w, h)
        mem_dc.SelectObject(bmp)

        # BitBlt 复制窗口内容
        mem_dc.BitBlt((0, 0), (w, h), src_dc, (left, top), win32con.SRCCOPY)

        # 转 numpy
        signed_bytes = bmp.GetBitmapBits(True)
        frame = np.frombuffer(signed_bytes, dtype=np.uint8)
        frame = frame.reshape((h, w, 4))  # BGRA
        frame = frame[:, :, :3]  # BGR

        # 清理资源
        src_dc.DeleteDC()
        mem_dc.DeleteDC()
        win32gui.ReleaseDC(self.hwnd, hwnd_dc)
        win32gui.DeleteObject(bmp.GetHandle())

        return frame
